"""
Moveable.py
"""
# import math
# import time
# import numpy as np
import sys
# import cubic_spline_planner
sys.path.append('..')
# from driver_model.IDM import IDM
from Dataloader.src import waymo_training_plot_tools as plot_tool
from Dataloader.src import waymo_scenario_classes as class_tool
from Dataloader.src.waymo_transfer_training_to_tf import *
from Dataloader.src.waymo_training_math_tools import *

'''
Moveable represents a general vehicle object with position, velocity eyc. In
 * each time step, the objects are updated by moving them forward (method
 * translate), by changing the velocity (method accelerate), and, in time steps
 * where timeToChange return true, by possibly changing the lane (method change) 
 * Implementations of this interface include
 *     Car: all "normal" cars and trucks
 *     Obstacle: These are implemented by immobile vehicles
 *     BCCar: These are introduced for bookkeeping purposes, only to make sure
 * that any Car has a predecessor required for calculating the acceleration.
'''
def measure_time(func):
    def wrapper(*args, **kwargs):
        start_time = time.time()
        result = func(*args, **kwargs)
        end_time = time.time()
        execution_time = end_time - start_time
        print(f"函数 {func.__name__} 运行时间：{execution_time:.6f} 秒")
        return result
    return wrapper
class lane(object):
    """
    lane车道类
    """
    def __init__(self,id,lastId,sideId,crossingId,width):
        self.id = id
        self.lastId = lastId
        self.sideId = sideId
        self.crossingId = crossingId
        self.width = width


class Moveable(object):
    def __init__(self,model, pos, vel, lane, length, width,acc=[0,0]):
        # lane需要用一个类来保存，这个类中有属性：车辆当前所在车道号id，以及左右车道号lid,rid(list sideId),小段前后车道号fid,bid(list crossingId)，车道宽度width
        self.model, self.pos, self.vel, self.lane, self.length, self.width = model, np.array(pos),np.array(vel), lane, length, width
        self.acc = np.array(acc)
        self.vel_array = []
        self.vel_array.append(vel)
        self.c_d = 0
        self.c_d_d = 0
        self.c_d_dd = 0
        self.Rc_d = 0
        self.Rc_d_d = 0
        self.Rc_d_dd = 0
        

    def calc_acc(self):
        """利用速度表最后0.3s内来计算加速度"""
        if(len(self.vel_array)<4):
            acc = [0,0,0]
            return
        acc_x = (self.vel_array[-1][0]-self.vel_array[-4][0])/0.3
        acc_y = (self.vel_array[-1][1]-self.vel_array[-4][1])/0.3
        acc = [acc_x,acc_y,math.sqrt(acc_x**2+acc_y**2)]
        self.acc = np.array(acc)

    def update_pos(self,pos):
        self.pos = np.array(pos)
    def update_d(self,c_d):
        self.c_d = c_d
    def update_d_d(self,c_d_d):
        self.c_d_d = c_d_d
    def update_d_dd(self,c_d_dd):
        self.c_d_dd = c_d_dd
    def update_vel(self,vel):
        self.vel = np.array(vel)
        self.vel_array.append(vel)

    def distance_to(self, fwd):
        """计算两车间间距

        Args:
            fwd (_type_): _description_

        Returns:
            _type_: _description_
        """
        lenth = np.array([self.length,self.length])
        distance = np.array(fwd.pos) - np.array(self.pos) - lenth
        return math.sqrt(distance[0]**2+distance[1]**2)


def moveableCreate(dataloader:Dataloader,timeIndex:int):
    result = {}
    initLane = lane(0,0,0,0,0)
    for id in dataloader.get_all_car_id_list():
        my_state_list = dataloader.get_agent_states_list_byId(id)
        my_state = my_state_list[0]
        pos = [my_state[3], my_state[4]]
        vel = [my_state[10], my_state[11], my_state[12]]
        object = Moveable(None,pos,vel,initLane,my_state[6],my_state[7])
        result[id] = object
    return result
        
        
@measure_time    
def moveableInitialize(dataloader,Id:int,object:Moveable,model):
    """车辆初始化

    Args:
        dataloader (_type_): 当前数据集的dataloader
        Id (_type_): 车辆的序号
        object (_type_): 车辆的Moveable对象

    Returns:
        _type_: _description_
    """
    my_state_list = dataloader.get_agent_states_list_byId(Id)
    my_state = my_state_list[0]
    pos = [my_state[3], my_state[4]]
    vel = [my_state[10], my_state[11], my_state[12]]
    a,_ = dataloader.get_agent_all_time_lane_list_byID(Id)
    id = a[0][0]    # 按照数据集接口中lane的状态位置注意调整
    sideId = [a[0][1],a[0][2]]  # 左道id，右道id
    crossingId = [a[0][4],a[0][3]]
    # print(sideId)
    if(not len(sideId[1])==0):
        # print("sideId",sideId[1][0])
        width,_ = dataloader.get_distance_to_around_lane_byID(id)
        # print("width:",width)
    else:
        width = -24
    # 0902change:delete the object.lane.id in initialize func
    myLane = lane(id,id,sideId,crossingId,width)
    length, width = my_state[6], my_state[7]
    object.model,object.pos,object.vel,object.lane,object.length,object.width = model,pos,vel,myLane,length,width
    return object

@measure_time
def moveableUpdate(dataloader,Id:int,time_index:int,object:Moveable,model,c_d_d=None,c_d_dd=None):
    """车辆更新

    Args:
        dataloader (_type_): 当前数据集的dataloader
        time_index (_type_): 当前时间步
        Id (_type_): 车辆的序号
        object (_type_): 车辆的Moveable对象

    Returns:
        _type_: _description_
    """
    # test.py    
    my_state_list = dataloader.get_agent_states_list_byId(Id)
    my_state = my_state_list[time_index]
    pos = [my_state[3], my_state[4]]
    vel = [my_state[10], my_state[11], my_state[12]]
    # left,right = dataloader.get_lane_boundary_byID(object.lane.lastId)  # 注意：在转弯途中是获取不到左道右道的
    a,_ = dataloader.get_agent_all_time_lane_list_byID(Id)
    id = a[time_index][0]    # 按照数据集接口中lane的状态位置注意调整
    sideId = [a[time_index][1],a[time_index][2]]
    crossingId = [a[time_index][4],a[time_index][3]]
    # print(sideId)
    # print(crossingId)
    if(not len(sideId[0])==0):
        # print("sideId",sideId[1][0])
        width,_ = dataloader.get_distance_to_around_lane_byID(id)
        # print("width:",width)
    else:
        width = []
    print("myLane's lastid",object.lane.id)
    myLane = lane(id,object.lane.id,sideId,crossingId,width)
    object.update_pos(pos)
    object.update_vel(vel)
    object.calc_acc()
    object.lane = myLane
    object.model = model
    return object
